import numpy as np
import gymnasium as gym
from gymnasium import spaces
from typing import Tuple, Optional, List

# Simple grid maze environment suitable for PPO with MLPPolicy.
# Observation: Box(shape=(4,)) with normalized [agent_x, agent_y, goal_x, goal_y].
# Action: Discrete(4) -> 0: up, 1: down, 2: left, 3: right.
# Reward: step -0.05, hit wall/obstacle -0.2, goal +1.0. Episode ends on goal or max_steps.

COLOR_FREE = np.array([255, 255, 255], dtype=np.uint8)  # white
COLOR_OBST = np.array([0, 0, 0], dtype=np.uint8)        # black
COLOR_AGENT = np.array([30, 144, 255], dtype=np.uint8)  # dodgerblue
COLOR_GOAL = np.array([50, 205, 50], dtype=np.uint8)    # limegreen


class MazeEnv(gym.Env):
    metadata = {"render_modes": ["rgb_array"], "render_fps": 10}

    def __init__(self, size: int = 10, obstacle_density: float = 0.18, max_steps: Optional[int] = None, render_mode: Optional[str] = None, seed: Optional[int] = 42):
        super().__init__()
        assert size >= 4, "Maze size must be >= 4"
        assert 0.0 <= obstacle_density <= 0.8, "Obstacle density should be in [0, 0.8]"
        self.size = size
        self.obstacle_density = obstacle_density
        self.max_steps = max_steps or (size * size)
        self.render_mode = render_mode
        self._rng = np.random.default_rng(seed)

        # Spaces
        self.observation_space = spaces.Box(low=0.0, high=1.0, shape=(4,), dtype=np.float32)
        self.action_space = spaces.Discrete(4)

        # Maze grid: 0 free, 1 obstacle
        self.grid = np.zeros((size, size), dtype=np.uint8)
        self.agent_pos = (0, 0)
        self.goal_pos = (size - 1, size - 1)
        self.steps = 0

        self._generate_maze()

    def _dfs_path(self, start: Tuple[int, int], goal: Tuple[int, int]) -> Optional[List[Tuple[int, int]]]:
        visited = set([start])
        parent: dict[Tuple[int, int], Tuple[int, int]] = {}
        stack: List[Tuple[int, int]] = [start]
        while stack:
            x, y = stack[-1]
            if (x, y) == goal:
                # reconstruct
                path: List[Tuple[int, int]] = []
                cur = goal
                path.append(cur)
                while cur != start:
                    cur = parent[cur]
                    path.append(cur)
                path.reverse()
                return path
            neighbors: List[Tuple[int, int]] = []
            for dx, dy in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
                nx, ny = x + dx, y + dy
                if 0 <= nx < self.size and 0 <= ny < self.size and (nx, ny) not in visited:
                    neighbors.append((nx, ny))
            if neighbors:
                self._rng.shuffle(neighbors)
                nxt = neighbors[0]
                parent[nxt] = (x, y)
                visited.add(nxt)
                stack.append(nxt)
            else:
                stack.pop()
        return None

    def _generate_path(self, start: Tuple[int, int], goal: Tuple[int, int]) -> List[Tuple[int, int]]:
        """Generate a complex path via multiple DFS segments through random detour waypoints.
        Allows up/left moves and produces a more winding route. Always ensures connectivity.
        """
        # choose detours based on size; cap to avoid extreme lengths
        detours = max(2, min(8, self.size // 2))
        # sample unique detour points excluding start/goal
        points = [start]
        chosen = set([start, goal])
        attempts = 0
        while len(points) - 1 < detours and attempts < detours * 10:
            i = int(self._rng.integers(0, self.size))
            j = int(self._rng.integers(0, self.size))
            if (i, j) not in chosen:
                chosen.add((i, j))
                points.append((i, j))
            attempts += 1
        points.append(goal)
        # connect sequentially via DFS; fallback to monotonic if a segment fails
        path: List[Tuple[int, int]] = []
        cur = points[0]
        path.append(cur)
        for idx in range(1, len(points)):
            nxt = points[idx]
            seg = self._dfs_path(cur, nxt)
            if seg is None:
                # fallback: monotonic
                x, y = cur
                gx, gy = nxt
                moves = np.array(['D'] * max(0, gx - x) + ['R'] * max(0, gy - y) + ['U'] * max(0, x - gx) + ['L'] * max(0, y - gy), dtype='<U1')
                if moves.size > 0:
                    self._rng.shuffle(moves)
                    for m in moves:
                        if m == 'D':
                            x += 1
                        elif m == 'U':
                            x -= 1
                        elif m == 'R':
                            y += 1
                        else:  # 'L'
                            y -= 1
                        seg_point = (x, y)
                        path.append(seg_point)
                cur = nxt
                continue
            # append seg, skipping duplicate of starting node
            for k in range(1, len(seg)):
                path.append(seg[k])
            cur = nxt
        # ensure final connectivity to goal (already last segment), return
        return path

    def _monotonic_path(self, start: Tuple[int, int], goal: Tuple[int, int]) -> List[Tuple[int, int]]:
        """Generate a minimal-length path from start to goal using axis-aligned moves.
        The order of moves is randomized to avoid a straight line but length is minimal.
        """
        x, y = start
        gx, gy = goal
        moves: List[Tuple[int, int]] = []
        dx = gx - x
        dy = gy - y
        if dx > 0:
            moves += [(1, 0)] * dx
        elif dx < 0:
            moves += [(-1, 0)] * (-dx)
        if dy > 0:
            moves += [(0, 1)] * dy
        elif dy < 0:
            moves += [(0, -1)] * (-dy)
        if moves:
            self._rng.shuffle(moves)
        path: List[Tuple[int, int]] = [start]
        for mx, my in moves:
            x += mx
            y += my
            path.append((x, y))
        return path

    def _generate_maze(self):
        # Path-first generation: build a guaranteed route, then fill obstacles on non-route cells by ratio
        self.grid[:, :] = 0
        start = (0, 0)
        goal = (self.size - 1, self.size - 1)
        path = self._generate_path(start, goal)
        path_set = set(path)

        # Compute target obstacles against total cells, evaluate path length budget
        total_cells = self.size * self.size
        desired_obstacles = int(round(self.obstacle_density * total_cells))
        max_free_cells = total_cells - desired_obstacles
        if len(path_set) > max_free_cells:
            eff_max_density = (total_cells - len(path_set)) / float(total_cells)
            print(f"[WARN] Path uses {len(path_set)} cells; max achievable density is {eff_max_density:.3f} (< desired {self.obstacle_density:.3f}). Falling back to monotonic path.")
            path = self._monotonic_path(start, goal)
            path_set = set(path)

        # Candidate cells for obstacles: all non-path cells
        candidates = [(i, j) for i in range(self.size) for j in range(self.size) if (i, j) not in path_set]

        target_obstacles = min(desired_obstacles, len(candidates))
        if target_obstacles > 0 and len(candidates) > 0:
            idx = self._rng.choice(len(candidates), size=target_obstacles, replace=False)
            for k in idx:
                ci, cj = candidates[k]
                self.grid[ci, cj] = 1

        # Enforce at least one obstacle per row/column without touching path cells
        # Rows
        for i in range(self.size):
            if np.sum(self.grid[i, :]) == 0:
                available = [(i, j) for j in range(self.size) if (i, j) not in path_set and self.grid[i, j] == 0]
                if available:
                    ci, cj = available[int(self._rng.integers(0, len(available)))]
                    self.grid[ci, cj] = 1
        # Columns
        for j in range(self.size):
            if np.sum(self.grid[:, j]) == 0:
                available = [(i, j) for i in range(self.size) if (i, j) not in path_set and self.grid[i, j] == 0]
                if available:
                    ci, cj = available[int(self._rng.integers(0, len(available)))]
                    self.grid[ci, cj] = 1

        # Ensure path cells are free (no global row/column clearing)
        for (i, j) in path_set:
            self.grid[i, j] = 0

        # Log achieved vs desired density
        achieved_obstacles = int(np.sum(self.grid))
        achieved_density = achieved_obstacles / float(total_cells)
        print(f"[INFO] Maze density: desired_obs={desired_obstacles}, placed_obs={achieved_obstacles}, achieved_density={achieved_density:.3f}, path_cells={len(path_set)}, candidates={len(candidates)}")

    def _norm_obs(self) -> np.ndarray:
        ax, ay = self.agent_pos
        gx, gy = self.goal_pos
        return np.array([
            ax / (self.size - 1),
            ay / (self.size - 1),
            gx / (self.size - 1),
            gy / (self.size - 1),
        ], dtype=np.float32)

    def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None):
        if seed is not None:
            self._rng = np.random.default_rng(seed)
        self.steps = 0
        # Start at top-left, goal at bottom-right
        self.agent_pos = (0, 0)
        self.goal_pos = (self.size - 1, self.size - 1)
        # Re-generate maze each reset only if requested
        regen = options.get("regen", False) if options else False
        if regen:
            self._generate_maze()
        obs = self._norm_obs()
        info = {}
        return obs, info

    def step(self, action: int):
        assert self.action_space.contains(action), "Invalid action"
        self.steps += 1
        reward = -0.05
        terminated = False
        truncated = False
        info = {}

        x, y = self.agent_pos
        nx, ny = x, y
        if action == 0:  # up
            nx = x - 1
        elif action == 1:  # down
            nx = x + 1
        elif action == 2:  # left
            ny = y - 1
        elif action == 3:  # right
            ny = y + 1

        # Bounds check
        if nx < 0 or nx >= self.size or ny < 0 or ny >= self.size:
            # hit wall
            reward -= 0.2
        elif self.grid[nx, ny] == 1:
            # hit obstacle
            reward -= 0.2
        else:
            self.agent_pos = (nx, ny)

        if self.agent_pos == self.goal_pos:
            reward += 1.0
            terminated = True

        if self.steps >= self.max_steps:
            truncated = True

        obs = self._norm_obs()
        return obs, reward, terminated, truncated, info

    def render(self):
        # Returns an RGB array representing the maze
        cell = 16
        H = self.size * cell
        W = self.size * cell
        img = np.tile(COLOR_FREE, (H, W, 1))

        # Draw obstacles
        for i in range(self.size):
            for j in range(self.size):
                if self.grid[i, j] == 1:
                    img[i*cell:(i+1)*cell, j*cell:(j+1)*cell] = COLOR_OBST

        # Draw goal
        gx, gy = self.goal_pos
        img[gx*cell:(gx+1)*cell, gy*cell:(gy+1)*cell] = COLOR_GOAL

        # Draw agent
        ax, ay = self.agent_pos
        img[ax*cell:(ax+1)*cell, ay*cell:(ay+1)*cell] = COLOR_AGENT

        return img

    def close(self):
        pass